import datetime
import time
import typing
from itertools import *

import joblib
import numpy as np
from scipy.stats import norm
from base_funcs import *


def init_base_model(_T=3, _v_0=1, _v_1=4, _v_2=6, _v_3=8, _p=19, _b=19, _c_e=10, _c_m=2, _h_e=1, _h_m=0.2):
    global alpha, K_t, p_K_t, D_t, p_D_t, p, b, c_e, c_m, h_e, h_m, theta
    global epoch, max_epoch, global_cache
    # Variable
    alpha = 0.95
    # Kt
    mu_K_t = 6  # 均值
    cv_K_t = 0.3  # 变异系数
    sigma_K_t = cv_K_t * mu_K_t  # 标准差
    K_t = np.array([4, 5, 6, 7, 8])
    p_K_t = p_norm(K_t, mu_K_t, sigma_K_t)
    # Dt
    mu_t = 5
    sigma_t = 0.5
    D_t = np.array(list(range(0, 10)))
    p_D_t = p_norm(D_t, mu_t, sigma_t)
    p = _p
    b = _b
    c_e = _c_e
    c_m = _c_m
    h_e = _h_e
    h_m = _h_m
    theta = h_m / (1 - alpha)
    # theta = 9, 109.8477
    # theta = 4, 109.8477
    print(f"c_m: {c_m}, h_e: {h_e}, h_m: {h_m}, theta: {theta}")
    joblib.dump({"T": _T,
                 "alpha": alpha,
                 "K_t": K_t, "p_K_t": p_K_t,
                 "D_t": D_t, "p_D_t": p_D_t,
                 "p": p, "b": b,
                 "c_e": c_e, "c_m": c_m,
                 "h_e": h_e, "h_m": h_m,
                 "theta": theta}, "params.pkl")
    print("params.pkl saved")

    v_0 = _v_0
    v_1 = _v_1
    v_2 = _v_2
    v_3 = _v_3

    calc_begin_time = time.perf_counter()

    from base_model_funcs import f_t

    res = f_t(v_0, v_1, v_2, v_3, 1)
    print(res)

    calc_end_time = time.perf_counter()
    print(f"计算完成, 耗时 {calc_end_time - calc_begin_time: .2f} 秒")
    return res


if __name__ == '__main__':
    print(init_base_model(_T=3, _v_0=1, _v_1=8, _v_2=12, _v_3=15))
    # T=3, 109.8477
    # T=4,
    # T=5,
